package com.mini.grpc.spring.client;

import com.mini.grpc.client.AbstractStub;
import com.mini.grpc.client.ClientInterceptor;
import com.mini.grpc.client.ManagedChannel;
import com.mini.grpc.client.ManagedChannelBuilder;
import com.mini.grpc.spring.annotation.GrpcClient;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.stereotype.Component;
import org.springframework.util.ReflectionUtils;

import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.logging.Logger;

/**
 * gRPC客户端注册表
 * 负责发现、创建和管理所有gRPC客户端
 * 
 * @author Mini-gRPC
 */
@Component
public class GrpcClientRegistry implements BeanPostProcessor, ApplicationContextAware {
    
    private static final Logger logger = Logger.getLogger(GrpcClientRegistry.class.getName());
    
    private ApplicationContext applicationContext;
    private final Map<String, ManagedChannel> channels = new ConcurrentHashMap<>();
    private final Map<String, Object> stubs = new ConcurrentHashMap<>();
    
    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }
    
    @Override
    public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
        Class<?> beanClass = bean.getClass();
        
        // 处理字段注入
        ReflectionUtils.doWithFields(beanClass, field -> {
            GrpcClient grpcClient = field.getAnnotation(GrpcClient.class);
            if (grpcClient != null && grpcClient.enabled()) {
                injectGrpcClient(bean, field, grpcClient);
            }
        });
        
        return bean;
    }
    
    /**
     * 注入gRPC客户端
     * 
     * @param bean 目标Bean
     * @param field 目标字段
     * @param grpcClient 注解信息
     */
    private void injectGrpcClient(Object bean, Field field, GrpcClient grpcClient) {
        try {
            Class<?> stubType = field.getType();
            String clientName = getClientName(grpcClient, stubType);
            
            // 获取或创建存根
            Object stub = getOrCreateStub(clientName, stubType, grpcClient);
            
            // 注入字段
            field.setAccessible(true);
            field.set(bean, stub);
            
            logger.info("Injected gRPC client: " + clientName + " into " + 
                       bean.getClass().getSimpleName() + "." + field.getName());
            
        } catch (Exception e) {
            logger.severe("Failed to inject gRPC client: " + e.getMessage());
            throw new RuntimeException("Failed to inject gRPC client", e);
        }
    }
    
    /**
     * 获取或创建存根
     * 
     * @param clientName 客户端名称
     * @param stubType 存根类型
     * @param grpcClient 注解信息
     * @return 存根实例
     */
    private Object getOrCreateStub(String clientName, Class<?> stubType, GrpcClient grpcClient) {
        return stubs.computeIfAbsent(clientName + ":" + stubType.getName(), key -> {
            try {
                // 获取或创建通道
                ManagedChannel channel = getOrCreateChannel(clientName, grpcClient);
                
                // 创建存根
                return createStub(stubType, channel, grpcClient);
                
            } catch (Exception e) {
                logger.severe("Failed to create gRPC stub " + stubType.getName() + ": " + e.getMessage());
                throw new RuntimeException("Failed to create gRPC stub " + stubType.getName(), e);
            }
        });
    }
    
    /**
     * 获取或创建通道
     * 
     * @param clientName 客户端名称
     * @param grpcClient 注解信息
     * @return 管理通道
     */
    private ManagedChannel getOrCreateChannel(String clientName, GrpcClient grpcClient) {
        return channels.computeIfAbsent(clientName, key -> {
            try {
                String address = grpcClient.address();
                if (address.isEmpty()) {
                    throw new IllegalArgumentException("gRPC client address is required for " + clientName);
                }
                
                // 解析地址
                String[] parts = address.split(":");
                if (parts.length != 2) {
                    throw new IllegalArgumentException("Invalid gRPC client address format: " + address);
                }
                
                String host = parts[0];
                int port = Integer.parseInt(parts[1]);
                
                // 创建通道构建器
                ManagedChannelBuilder<?> channelBuilder = ManagedChannelBuilder.forAddress(host, port)
                        .usePlaintext(); // 默认使用明文传输
                
                // 配置通道参数
                configureChannel(channelBuilder, grpcClient);
                
                ManagedChannel channel = channelBuilder.build();
                logger.info("Created gRPC channel for " + clientName + " -> " + address);
                
                return channel;
                
            } catch (Exception e) {
                logger.severe("Failed to create gRPC channel for " + clientName + ": " + e.getMessage());
                throw new RuntimeException("Failed to create gRPC channel for " + clientName, e);
            }
        });
    }
    
    /**
     * 配置通道参数
     * 
     * @param channelBuilder 通道构建器
     * @param grpcClient 注解信息
     */
    private void configureChannel(ManagedChannelBuilder<?> channelBuilder, GrpcClient grpcClient) {
        // 设置负载均衡策略
        channelBuilder.defaultLoadBalancingPolicy(grpcClient.loadBalancingPolicy());
        
        // 设置keepalive
        if (grpcClient.keepAliveEnabled()) {
            channelBuilder.keepAliveTime(grpcClient.keepAliveTimeSeconds(), TimeUnit.SECONDS)
                         .keepAliveTimeout(grpcClient.keepAliveTimeoutSeconds(), TimeUnit.SECONDS)
                         .keepAliveWithoutCalls(grpcClient.keepAliveWithoutCalls());
        }
        
        // 设置最大消息大小
        channelBuilder.maxInboundMessageSize(4 * 1024 * 1024); // 4MB默认值
    }
    
    /**
     * 创建存根
     * 
     * @param stubType 存根类型
     * @param channel 通道
     * @param grpcClient 注解信息
     * @return 存根实例
     */
    private Object createStub(Class<?> stubType, ManagedChannel channel, GrpcClient grpcClient) throws Exception {
        // 查找存根工厂方法
        Method factoryMethod = findStubFactoryMethod(stubType);
        if (factoryMethod == null) {
            throw new IllegalArgumentException("No factory method found for stub type: " + stubType.getName());
        }
        
        // 创建存根
        Object stub = factoryMethod.invoke(null, channel);
        
        // 应用拦截器
        if (grpcClient.interceptors().length > 0) {
            stub = applyInterceptors(stub, grpcClient.interceptors());
        }
        
        // 设置调用选项
        if (stub instanceof AbstractStub) {
            AbstractStub<?> abstractStub = (AbstractStub<?>) stub;
            if (grpcClient.callTimeoutMillis() > 0) {
                stub = abstractStub.withDeadlineAfter(grpcClient.callTimeoutMillis(), TimeUnit.MILLISECONDS);
            }
        }
        
        return stub;
    }
    
    /**
     * 查找存根工厂方法
     * 
     * @param stubType 存根类型
     * @return 工厂方法
     */
    private Method findStubFactoryMethod(Class<?> stubType) {
        // 查找 newBlockingStub, newStub, newFutureStub 等方法
        String[] factoryMethodNames = {"newBlockingStub", "newStub", "newFutureStub"};
        
        for (String methodName : factoryMethodNames) {
            try {
                Method method = stubType.getMethod(methodName, ManagedChannel.class);
                if (method.getReturnType().isAssignableFrom(stubType)) {
                    return method;
                }
            } catch (NoSuchMethodException e) {
                // 继续查找下一个方法
            }
        }
        
        return null;
    }
    
    /**
     * 应用拦截器
     * 
     * @param stub 存根
     * @param interceptorClasses 拦截器类数组
     * @return 应用拦截器后的存根
     */
    private Object applyInterceptors(Object stub, Class<?>[] interceptorClasses) {
        if (!(stub instanceof AbstractStub)) {
            return stub;
        }
        
        AbstractStub<?> abstractStub = (AbstractStub<?>) stub;
        
        for (Class<?> interceptorClass : interceptorClasses) {
            try {
                Object interceptorBean = applicationContext.getBean(interceptorClass);
                if (interceptorBean instanceof ClientInterceptor) {
                    abstractStub = abstractStub.withInterceptors((ClientInterceptor) interceptorBean);
                } else {
                    logger.warning("Interceptor class " + interceptorClass.getName() + 
                                 " does not implement ClientInterceptor");
                }
            } catch (Exception e) {
                logger.warning("Failed to resolve interceptor " + interceptorClass.getName() + ": " + e.getMessage());
            }
        }
        
        return abstractStub;
    }
    
    /**
     * 获取客户端名称
     * 
     * @param grpcClient 注解信息
     * @param stubType 存根类型
     * @return 客户端名称
     */
    private String getClientName(GrpcClient grpcClient, Class<?> stubType) {
        String name = grpcClient.value();
        if (name.isEmpty()) {
            name = stubType.getSimpleName();
        }
        return name;
    }
    
    /**
     * 关闭所有通道
     */
    public void shutdown() {
        logger.info("Shutting down " + channels.size() + " gRPC channels");
        
        for (Map.Entry<String, ManagedChannel> entry : channels.entrySet()) {
            try {
                ManagedChannel channel = entry.getValue();
                channel.shutdown();
                
                if (!channel.awaitTermination(5, TimeUnit.SECONDS)) {
                    logger.warning("Channel " + entry.getKey() + " did not terminate gracefully, forcing shutdown");
                    channel.shutdownNow();
                }
                
            } catch (Exception e) {
                logger.warning("Error shutting down channel " + entry.getKey() + ": " + e.getMessage());
            }
        }
        
        channels.clear();
        stubs.clear();
    }
    
    /**
     * 获取通道数量
     * 
     * @return 通道数量
     */
    public int getChannelCount() {
        return channels.size();
    }
    
    /**
     * 获取存根数量
     * 
     * @return 存根数量
     */
    public int getStubCount() {
        return stubs.size();
    }
} 